'''
多元线性回归, 训练和测试切分
'''

import xlrd
import numpy as np
import matplotlib.pyplot as plt
from pylab import *
mpl.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] =False
from sklearn.linear_model import LinearRegression
import datetime
from xlrd import xldate_as_tuple
from sklearn.ensemble import AdaBoostClassifier
from sklearn.ensemble import GradientBoostingRegressor
import time

file_path = '22个省市_平均价_猪肉.xlsx'

work_book = xlrd.open_workbook(file_path)
content = work_book.sheet_by_name('22个省市_平均价_猪肉')

time_list = content.col_values(0)[1:]
pig_price_list = content.col_values(1)[1:]    # 猪肉价格
profit_list = content.col_values(2)[1:]       # 养殖利润
import_num_list = content.col_values(3)[1:]   # 进口数量
invemtory_list = content.col_values(4)[1:]    # 生猪存栏，能繁母猪

output_price_list = []
output_profit_list =[]
output_import_num_list = []
output_invemtory_list = []
output_time_list = []

new_price = 0
new_profit = 0
new_import_num = 0
new_invemtory = 0

for the_time, pig_price, profit, import_num, invemtory in zip(time_list, pig_price_list, profit_list, import_num_list, invemtory_list):
    if pig_price != '':
        new_price = pig_price
    if profit != '':
        new_profit = profit
    if import_num != '':
        new_import_num = import_num
    if invemtory != '':
        new_invemtory = invemtory
    if import_num != '' or invemtory != '':
        output_import_num_list.append(new_import_num)
        output_invemtory_list.append(new_invemtory)
        output_profit_list.append(new_profit)
        output_price_list.append(new_price)
        output_time_list.append(the_time)

input_time_list = []
for the_time in output_time_list:
    new_date = datetime.datetime(*xldate_as_tuple(the_time, 0))
    input_time_list.append(new_date)
# time.sleep(500)

gap_len = 10  # 预测的间隔时间

add_datetime_list = ['2022-8-31', '2022-9-30', '2022-10-31', '2022-11-30', '2022-12-31',
                     '2023-1-31', '2023-2-28', '2023-3-31', '2023-4-30', '2022-5-31']

input_time_list = input_time_list[gap_len:]
for date_str in add_datetime_list:
    fmt = '%Y-%m-%d'
    time_tuple = time.strptime(date_str, fmt)
    year, month, day = time_tuple[:3]
    a_date = datetime.date(year, month, day)
    print(a_date)
    input_time_list.append(datetime.date(*map(int, date_str.split('-'))))
print(input_time_list)


input_price_list = output_price_list[gap_len:]
input_profit_list = output_profit_list[:]          # 对未知实际价格数据进行未来预测
input_import_num_list = output_import_num_list[:]
input_invemtory_list = output_invemtory_list[:]

print(len(input_profit_list), len(input_profit_list), len(input_import_num_list), len(input_invemtory_list))

# 绘制散点图
# plt.scatter(input_price_list[350:], input_profit_list[350:])
# plt.show()


# ada = AdaBoostClassifier(n_estimators=180, random_state=0)
# gb = GradientBoostingRegressor(max_depth=4, n_estimators=200, random_state=2)

X = input_profit_list
A = input_import_num_list
B = input_invemtory_list

Y = input_price_list

x = np.array(X)
a = np.array(A)
b = np.array(B)
y = np.array(Y)
X = np.c_[np.ones(len(x)), a, b]  # 将三个特征组合输入
Y = y.astype('int')

#把X,Y放入模型进行训练
lr = LinearRegression()
lr.fit(X[:60], Y[:60])

#对X数据进行预测
pre_y = lr.predict(X[60:])
# y = Y[60:] + ['' for i in range(gap_len)]
x = input_time_list[60:]
fig, ax = plt.subplots(figsize=(6, 6))
ax.xaxis.grid(True)
ax.yaxis.grid(True)
ax.plot(x[:-gap_len], Y[60:], c='blue', label='实际价格')
ax.plot(x, pre_y, 'r', label='预测价格')
fig.autofmt_xdate()
plt.legend()
plt.tight_layout()
plt.show()

